package com.example.qwen02;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.tools.FunctionDefinition;
import com.alibaba.dashscope.tools.ToolCallBase;
import com.alibaba.dashscope.tools.ToolFunction;
import com.alibaba.dashscope.utils.JsonUtils;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.victools.jsonschema.generator.*;
import com.alibaba.dashscope.tools.ToolCallFunction;

import java.util.*;

public class Fun01 {
    public static class WeekTool{
        public  String weekToday(){
            String[] names=new String[]{"星期日","星期一","星期二","星期三","星期四","星期五","星期六"};
            Calendar c=Calendar.getInstance();
            int w=c.get(Calendar.DAY_OF_WEEK);
            return "今天是"+names[w-1];
        }
    }
    public static void main(String[] args) throws NoApiKeyException, InputRequiredException {
        //创建构建，固定不变的
        SchemaGeneratorConfigBuilder configBuilder =
                new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON);
        SchemaGeneratorConfig config = configBuilder.with(Option.EXTRA_OPEN_API_FORMAT_VALUES)
                .without(Option.FLATTENED_ENUMS_FROM_TOSTRING).build();
        SchemaGenerator generator = new SchemaGenerator(config);

        ObjectNode jsonSchema_week = generator.generateSchema(WeekTool.class);
        FunctionDefinition fd_week = FunctionDefinition.builder().name("get_current_week").description("获取今天星期几")
                .parameters(JsonUtils.parseString(jsonSchema_week.toString()).getAsJsonObject()).build();

        Message systemMsg = Message.builder().role(Role.SYSTEM.getValue())
                .content("你是一个小助手，回答用户的问题，当遇到无法解决的问题的时候，可以调用工具方法")
                .build();
        Message userMsg = Message.builder().role(Role.USER.getValue())
                .content("昨天星期几")
                .build();
        List<Message> msgList=new LinkedList<>();
        msgList.add(systemMsg);
        msgList.add(userMsg);

        Generation gen = new Generation();
        QwenParam params = QwenParam.builder().model("qwen-plus")
                .tools(Arrays.asList(ToolFunction.builder().function(fd_week).build()))
                .messages(msgList)
                //.prompt("今天星期几?")
                .seed(222)
                .topP(0.8)
                .resultFormat("message")
                .enableSearch(false)
                .maxTokens(100)
                .temperature((float)0.85)
                .repetitionPenalty((float)1.0)
                .build();

        GenerationResult result = gen.call(params);
        System.out.println(result);

        //对结果进行分类处理
        for (GenerationOutput.Choice choice : result.getOutput().getChoices()) {
            msgList.add(choice.getMessage());//大模型返回的消息记录到上下文
            // 如果需要调用工具
            if (result.getOutput().getChoices().get(0).getMessage().getToolCalls() != null) {
                for (ToolCallBase toolCall : result.getOutput().getChoices().get(0).getMessage()
                        .getToolCalls()) {//有可能要求助调用多个方法
                    if (toolCall.getType().equals("function")) {
                        // 获取工具函数名称和入参
                        String functionName = ((ToolCallFunction) toolCall).getFunction().getName();
                        String functionArgument = ((ToolCallFunction) toolCall).getFunction().getArguments();
                        // 大模型判断调用天气查询工具的情况
                        if (functionName.equals("get_current_week")) {
                            //帮LLM的忙，去调用自己的方法
                            WeekTool weekTool=new WeekTool();
                           String todayWeek=weekTool.weekToday();

                            Message toolResultMessage = Message.builder().role("tool")
                                    .content(String.valueOf(todayWeek)).toolCallId(toolCall.getId()).build();
                            msgList.add(toolResultMessage);
                            System.out.println("\n工具输出信息："+todayWeek);
                        }
                    }
                }
            }
            // 如果无需调用工具，直接输出大模型的回复
            else {
                System.out.println("\n最终答案："+result.getOutput().getChoices().get(0).getMessage().getContent());
                return;
            }
        }
        // 大模型的第二轮调用 包含工具输出信息
        params.setMessages(msgList);
        result = gen.call(params);
        System.out.println("\n大模型第二轮输出信息："+JsonUtils.toJson(result));
        System.out.println(("\n最终答案："+result.getOutput().getChoices().get(0).getMessage().getContent()));

    }

}